{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "70143657",
   "metadata": {},
   "outputs": [],
   "source": [
    "from itertools import count  #用来计数\n",
    "import torch\n",
    "import torch.autograd\n",
    "import torch.nn.functional as F\n",
    "n = 3   #多项式最高次数为3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "238b7d6a",
   "metadata": {},
   "outputs": [],
   "source": [
    "#造数据\n",
    "def make_features(x):\n",
    "    \"\"\"构造特征矩阵\"\"\"\n",
    "    x = x.unsqueeze(1)  #把x中每个元素加一个中括号\n",
    "    return torch.cat([x**i for i in range(1,n+1)],1)  #实现tensor的拼接"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3556dba7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3\n"
     ]
    }
   ],
   "source": [
    "#定义多项式\n",
    "W_target = torch.FloatTensor([2,3,4]).unsqueeze(1)   #增加维度，给每个元素加个中括号\n",
    "b_target = torch.FloatTensor([1])\n",
    "print(W_target.size(0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "31eb9e93",
   "metadata": {},
   "outputs": [],
   "source": [
    "#根据输入的x得到一个拟合结果\n",
    "def f(x):\n",
    "    \"\"\"拟合函数\"\"\"\n",
    "    return x.mm(W_target) + b_target.item()    #x.mm表示做矩阵乘法"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "84cecba2",
   "metadata": {},
   "outputs": [],
   "source": [
    "#生成训练集（随机数为参数）\n",
    "#每次取batch_size个数据点，将其处理为矩阵形式，再传给f(x)\n",
    "def get_batch(batch_size = 32):\n",
    "    \"\"\"构建批次对(x,f(x))\"\"\"\n",
    "    random = torch.randn(batch_size)  #生成batch_size行，1列的随机数（符合正态分布）\n",
    "    x = make_features(random)\n",
    "    y = f(x)\n",
    "    return x,y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "0510703f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#定义模型\n",
    "#输入值和目标参数w的行数一致，输出值为1的模型\n",
    "fc = torch.nn.Linear(W_target.size(0),1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "36bf9579",
   "metadata": {},
   "outputs": [],
   "source": [
    "#训练模型直到均方差小于0.001\n",
    "for batch_idx in count(1):\n",
    "    #获取数据\n",
    "    batch_x, batch_y = get_batch()  #batch_x为特征值随机的特征矩阵，batch_y为该特征矩阵生成的拟合结果\n",
    "    #垂直梯度\n",
    "    fc.zero_grad()\n",
    "    #前向传播\n",
    "    output = F.smooth_l1_loss(fc(batch_x),batch_y)   #fc(batch_x)是训练结果，batch_y为实际结果\n",
    "    loss = output.item()\n",
    "    #反向传播\n",
    "    output.backward()\n",
    "    #计算梯度\n",
    "    for param in fc.parameters():\n",
    "        param.data.add_(-0.1*param.grad.data)\n",
    "    #设置截至条件\n",
    "    if loss < 1e-3:\n",
    "        break;\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e4b17dc9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "loss:0.000926 after 98 batches\n",
      "==>Learned function:\ty = +2.00x^3+2.97x^2+4.00x^1+1.05\n",
      "==>Actual function:\ty = +2.00x^3+3.00x^2+4.00x^1+1.00\n"
     ]
    }
   ],
   "source": [
    "#打印函数\n",
    "def poly_desc(W,b):\n",
    "    \"\"\"建立多项式线性描述\"\"\"\n",
    "    result = 'y = '\n",
    "    for i,w in enumerate(W):\n",
    "        result += '{:+.2f}x^{}'.format(w,len(W)-i)\n",
    "    result += '{:+.2f}'.format(b[0])\n",
    "    return result\n",
    "\n",
    "print('loss:{:.6f} after {} batches'.format(loss,batch_idx))\n",
    "print('==>Learned function:\\t'+poly_desc(fc.weight.view(-1),fc.bias))\n",
    "print('==>Actual function:\\t'+poly_desc(W_target.view(-1),b_target))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dea7b624",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "language": "python",
   "name": "pytorch"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
